'''
Created on Aug 19, 2009

@author: mkiyer
'''
from bx.intervals.intersection import Interval, IntervalTree
from interval_io import parse_raw_interval_file, parse_bed_file
from plots import plot_bar

import glob
import os
import collections
import logging
import sys

def load_intervals(dbpath):
    trees = collections.defaultdict(lambda: IntervalTree())
    intervals = []

    files = glob.glob(os.path.join(dbpath, '*.bed'))
    for f in files:
        logging.debug("loading intervals from file %s" % f)
        category_name = os.path.splitext(os.path.basename(f))[0]
        for interval in parse_raw_interval_file(open(f),
                                                category_name):
            trees[interval.chrom].insert_interval(interval)
            intervals.append(interval)
    return trees, intervals

if __name__ == '__main__':
    '''
    usage: python <prog> inputfile outputfile
    '''
    inputfile = sys.argv[1]
    outputfile = sys.argv[2]
    
    inputcat = os.path.splitext(os.path.basename(inputfile))[0]
    testintervals = [x for x in parse_raw_interval_file(open(inputfile), inputcat)]

    logging.basicConfig(level=logging.DEBUG)
    dbpath = '/home/mkiyer/exper/ncrna/classify'
    trees, intervals = load_intervals(dbpath)

    categories = {'coding': frozenset(['ucsc_coding_exons']),
                  'noncoding': frozenset(['ucsc_noncoding_transcripts',
                                          'ucsc_coding_introns',
                                          'lincrna_exons',
                                          'perez_long_ncrna_cancer',
                                          'PINS',
                                          'TINS',
                                          'rnadb_antisense',
                                          'rnadb_mirna'])}
    # TODO: hack -- just finds unannotated intervals for now
    desired_bin = 2
    
    cat_bin_pos = {}
    cat_bin_names = {}
    for i, k in enumerate(categories.keys()):
        cat_bin_pos[k] = i
        cat_bin_names[i] = k
    # group intervals accord to category definitions        
    cat_bins = collections.defaultdict(lambda: 0)
    cat_counts = collections.defaultdict(lambda: 0)
    
    outfhd = open(outputfile, 'w')
    
    for interval in testintervals:
        hits = trees[interval.chrom].find(interval.start, interval.end)
        cat_hits = {}
        bin = 0
        # assign overlapping hits to categories
        for hit in hits:
            for cat, members in categories.iteritems():
                if hit.value[1] in members:
                    cat_hits[cat] = 1
                    bin |= (1 << cat_bin_pos[cat])
        # count numbers in each category
        for k, v in cat_hits.iteritems():
            cat_counts[k] += v
        # bin according to combinations of different categories
        cat_bins[bin] += 1
        # if this hit matches desired properties, then write to file
        # TODO: hack!
        if bin == desired_bin:
            outfhd.write('%s\t%d\t%d\t%s\n' %
                         (interval.chrom, interval.start, 
                          interval.end, interval.value[0]))

    outfhd.close()
    
    # construct bin names -> bin counts
    labels = []
    bars = []
    for bin, count in cat_bins.iteritems():    
        # reconstruct bin name
        bin_name = []    
        for i in xrange(len(categories.keys())):
            if bin & (1 << i) != 0:
                bin_name.append(cat_bin_names[i])
        bin_name = '+'.join(bin_name)
        if len(bin_name) == 0:
            bin_name = 'none'
        labels.append(bin_name)
        bars.append(count)
        logging.debug("%s -> %s" % (bin_name, count))
    
    # plots
    plot_bar(labels, bars, outfile="category_counts.png",
             title="Interval classification",
             xlabel="Interval type",
             ylabel="Count")

    # output category counts
    print len(intervals), 'total intervals in db'
    print len(testintervals), 'total intervals in input file'
    print cat_counts
    
    
    